#' @TODO 高低风险组免疫浸润分析
#' @title ## 高低风险组免疫浸润分析
#' @param saveplot 是否保存图片
#' @param output_dir 文件保存路径，需要以`/`结尾
#' @param exp  表达矩阵，为了计算模型基因与免疫评分的相关性
#' @param model_gene lasso后 纳入模型的基因
#' @param risk_infor KM_ROC_curve计算结果第三个。或者一个包含`sample`与`riskscore`，`riskgroup`的*data.frame*
#' @param immucell_res 免疫浸润计算结果
#' 
#' @return 第一个元素是转化后的长表，第二个是figure
#' 
#' @author *WYK*
#'
Immune_infiltration <- function(immucell_res = NULL, saveplot = F, exp = NULL, model_gene = NULL,
    var_name = NULL, risk_infor = NULL, output_dir = "./", w = 8, h = 5) {
    library(tidyverse)
    library(ggpubr)
    library(cowplot)
    require(rstatix)

    if (is.null(var_name)) {
        var_name <- deparse(substitute(risk_infor))
    }
    CIBERSORT <- immucell_res

    CIBERSORT <- CIBERSORT %>%
        as.data.frame() %>%
        dplyr::rename(sample = 1)

    # # View(CIBERSORT) immu_res$cibersort

    CIBERSORT <- inner_join(CIBERSORT, risk_infor %>%
        dplyr::select(sample, riskgroup))


    CIBERSORT_long <- CIBERSORT %>%
        pivot_longer(cols = -c(sample, riskgroup), names_to = "type", values_to = "value") %>%
        filter(type != "P-value") %>%
        filter(type != "Correlation") %>%
        filter(type != "RMSE")

    d <- CIBERSORT_long %>% group_by(type) %>% summarise(summ = sum(value))
    
    if(length(which(d[[2]] == 0))>= 1){
        
        cli::cli_alert_warning("{d[[1]][which(d[[2]] == 0)]} 浸润程度等于0，去除")
        CIBERSORT_long %<>% filter(type != d[[1]][which(d[[2]] == 0)])
    }

    p_df <- CIBERSORT_long %>%
        group_by(type) %>%
        rstatix::wilcox_test(value ~ riskgroup) %>%
        rstatix::add_significance() %>%
        rstatix::add_xy_position(x = "type")

    figure_immune_infiltration <- ggplot(data = CIBERSORT_long, mapping = aes(
        x = type,
        y = value
    )) +
        geom_boxplot(aes(fill = riskgroup),
            size = 0.5, notch = F,
            outlier.size = 0.5
        ) +
        theme_pubr() +
        theme(axis.text.x = element_text(
            angle = 45,
            hjust = 1, vjust = 1
        )) +
        guides(fill = guide_legend(title = "Score")) +
        theme(legend.position = "top") +
        labs(title = "", x = "", y = "Fraction") +
        scale_fill_manual(values = c(Low = "#377EB8", High = "#E41A1C")) +
        theme(plot.margin = unit(c(
            0,
            0, 0, 1
        ), "cm"), axis.text = element_text(color = "black")) +
        ggpubr::stat_pvalue_manual(p_df,
            label = "p.signif", label.size = 3, bracket.size = 0, tip.length = 0.02
        )

    if (saveplot) {
        if (!dir.exists(sprintf("%s", output_dir))) {
            dir.create(sprintf("%s", output_dir), recursive = T)
        }

        ggsave2(filename = sprintf("%sFigure_%s_immu_infl_diff.pdf", output_dir,
            var_name), plot = figure_immune_infiltration, width = w, height = h)
    }

    tmp <- list()
    tmp[[1]] <- CIBERSORT_long
    tmp[[2]] <- figure_immune_infiltration
    names(tmp) <- c("CIBERSORT_long", "figure_immune_infiltration")

    CIBERSORT <- CIBERSORT_long %>%
        pivot_wider(names_from = type, values_from = value) %>%
        dplyr::select(-riskgroup) %>%
        as.data.frame() %>%
        column_to_rownames(var = "sample")

    if (!is.null(model_gene)) {
        fpkm_use_3 <- exp[model_gene, ] %>%
            t() %>%
            as.data.frame()

        fpkm_4 <- fpkm_use_3[rownames(CIBERSORT), ]

        fpkm_4 <- fpkm_4 %>%
            rownames_to_column(var = "sample") %>%
            left_join(., (risk_infor %>%
                dplyr::select(sample, riskscore))) %>%
            dplyr::rename(Score = riskscore)

        fpkm_4 <- fpkm_4 %>%
            column_to_rownames(var = "sample")

        cor_res <- psych::corr.test(x = as.matrix(fpkm_4), y = as.matrix(CIBERSORT),
            method = "spearman", adjust = "BH")
        cor_res$r %>%
            head()

        p <- pheatmap::pheatmap(as.matrix(cor_res$r), border = FALSE, main = "",
            display_numbers = matrix(ifelse(cor_res$p <= 0.05, ifelse(cor_res$p <=
                0.01, ifelse(cor_res$p <= 0.001, ifelse(cor_res$p <= 1e-04, "****",
                "***"), "**"), "*"), ""), nrow(cor_res$p)), color = colorRampPalette(c("#377EB8",
                "white", "firebrick3"))(50), cluster_rows = F, cluster_cols = F,
            fontsize_number = 11, angle_col = 45, border_color = "white")

        write.table(cor_res$r, file = sprintf("%scibersort_gene_cell_rho.txt", output_dir),
            row.names = TRUE, quote = F, sep = "\t")
        write.table(cor_res$p, file = sprintf("%scibersort_gene_cell_rho_p.txt",
            output_dir), row.names = TRUE, quote = F, sep = "\t")
        write.table(cor_res$p.adj, file = sprintf("%scibersort_gene_cell_rho_BH_adjp.txt",
            output_dir), row.names = TRUE, quote = F, sep = "\t")

        p <- ggplotify::as.ggplot(p)

        p <- p + theme(plot.margin = unit(c(0, 0, 0, 1.1), "cm"))

        anno_col <- data.frame(row.names = CIBERSORT_long %>%
            pivot_wider(names_from = type, values_from = value) %>%
            dplyr::select(sample) %>%
            pull(), RiskGroup = CIBERSORT_long %>%
            pivot_wider(names_from = type, values_from = value) %>%
            dplyr::select(riskgroup) %>%
            pull()) %>%
            arrange(RiskGroup)

        anno_colors <- list(RiskGroup = c(High = "#E41A1C", Low = "#367EB8"))

        p2 <- pheatmap::pheatmap(mat = t(CIBERSORT)[, rownames(anno_col)], scale = "row",
            cluster_rows = F, cluster_cols = F, show_colnames = F, annotation_col = anno_col,
            annotation_colors = anno_colors, color = colorRampPalette(c("#377EB8",
                "white", "firebrick3"))(50))

        p2 <- ggplotify::as.ggplot(p2)

        if (saveplot) {
            ggsave2(filename = sprintf("%sgene_cell_cor.pdf", output_dir), plot = p,
                width = 7, height = nrow(cor_res$r) * 0.5)


            ggsave2(filename = sprintf("%sheatmap_cell_sample.pdf", output_dir),
                plot = p2, width = 7, height = 4.5)
        }
    }

    return(tmp)
}

# (ns: p > 0.05，*: p <= 0.05，**: p <= 0.01， ***: p <= 0.001，****: p <=
# 0.0001)
